import cv2
import os
import numpy as np
import PIL.Image as Image
import pickle
import matplotlib.pyplot as plt
from sklearn.svm import LinearSVC
from sklearn.metrics import roc_curve, auc


def getTrainHoGsAndLabels(hog, posTrainImagesPath, negTrainImagesPath):
    dataList = []
    labelsList = []
    for image in posTrainImagesPath:
        # img = Image.open(image)
        # imgArr = np.array(img)
        # # print(imgArr.shape)
        # imgArr_norm = np.array(imgArr[16:144, 16:80])
        # print(imgArr_norm.shape)
        img_norm = cv2.imread(image)[16:144, 16:80]
        winStride = (8, 8)
        padding = (0, 0)
        hogValues = hog.compute(img_norm, winStride, padding)
        dataList.append(list(hogValues))
        labelsList.append([1])
        # print(hogValues.shape)

    for image in negTrainImagesPath:
        # img = Image.open(image)
        # imgArr = np.array(img)
        # # print(imgArr.shape)
        # imgArr_norm = np.array(imgArr[16:144, 16:80])
        # print(imgArr_norm.shape)
        
        img_norm = cv2.imread(image)
        imgShape = img_norm.shape
        # print(imgShape)
        for i in range(10):
            h = np.random.randint(0, imgShape[0] - 129)
            w = np.random.randint(0, imgShape[1] - 65)
            winStride = (8, 8)
            padding = (0, 0)
            hogValues = hog.compute(img_norm[h:h+128, w:w+64], winStride, padding)
            dataList.append(list(hogValues))
            labelsList.append([0])

    data = np.array(dataList).reshape((-1, 3780))
    # print(data.shape)
    labels = np.array(labelsList)
    # print(labels.shape)
    return data, labels

def getAccuracy(svm, hog, posTestImagesPath, negTestImagesPath):
    T = 0
    F = 0
    for image in posTestImagesPath:
        img_norm = cv2.imread(image)[3:131, 3:67]
        
        winStride = (8, 8)
        padding = (0, 0)
        hogValues = hog.compute(img_norm, winStride, padding).reshape((1, 3780))
        # print(hogValues)
        # print(type(hogValues[0][0]))
        # hogValuesF32 = np.array(hogValues, dtype="float32")
        pred = svm.predict(hogValues)[0]
        if pred == 1:
            T += 1
        else:
            F += 1

    for image in negTestImagesPath:
        img_norm = cv2.imread(image)[3:131, 3:67]
        
        winStride = (8, 8)
        padding = (0, 0)
        hogValues = hog.compute(img_norm, winStride, padding).reshape((1, 3780))
        # print(hogValues)
        # print(type(hogValues[0][0]))
        # hogValuesF32 = np.array(hogValues, dtype="float32")
        pred = svm.predict(hogValues)[0]
        if pred == 0:
            T += 1
        else:
            F += 1
    Acc = T / (T + F)
    return Acc

def getTestHoGsAndLabels(hog, posTestImagesPath, negTestImagesPath):
    testDataList = []
    testLabelsList = []
    for image in posTestImagesPath:
        # img = Image.open(image)
        # imgArr = np.array(img)
        # # print(imgArr.shape)
        # imgArr_norm = np.array(imgArr[16:144, 16:80])
        # print(imgArr_norm.shape)
        img_norm = cv2.imread(image)[3:131, 3:67]
        winStride = (8, 8)
        padding = (0, 0)
        hogValues = hog.compute(img_norm, winStride, padding)
        testDataList.append(list(hogValues))
        testLabelsList.append([1])
        # print(hogValues.shape)

    for image in negTestImagesPath:
        # img = Image.open(image)
        # imgArr = np.array(img)
        # # print(imgArr.shape)
        # imgArr_norm = np.array(imgArr[16:144, 16:80])
        # print(imgArr_norm.shape)
        img_norm = cv2.imread(image)[3:131, 3:67]
        winStride = (8, 8)
        padding = (0, 0)
        hogValues = hog.compute(img_norm, winStride, padding)
        testDataList.append(list(hogValues))
        testLabelsList.append([0])
        # print(hogValues.shape)
        
    testData = np.array(testDataList).reshape((-1, 3780))
    testLabels = np.array(testLabelsList)
    return testData, testLabels

def drawROCAndComputeAUC(svm, testData, testLabels):
    Y_score = svm.decision_function(testData)
    roc_auc = dict()
    fpr, tpr, threshold = roc_curve(testLabels, Y_score)
    roc_auc = auc(fpr, tpr)

    plt.figure()
    lw = 3
    plt.plot(fpr, tpr, color='darkorange',
                lw=lw, label='ROC curve (area = %0.2f)' % roc_auc)
    plt.plot([0, 1], [0, 1], color='navy', lw=lw, linestyle='--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('Receiver operating characteristic example')
    plt.legend(loc="lower right")
    plt.savefig("ROC.png")
    plt.show()
    AUC = 0.0
    for i in range(len(fpr)):
        if i == 0:
            AUC += tpr[i] * fpr[i]
        else:
            AUC += tpr[i] * (fpr[i] -fpr[i-1])
    return AUC

if __name__ == "__main__":
    posTrainFolderPath = ".\\INRIAPerson\\96X160H96\\Train\\posNew"
    posImgs = os.listdir(posTrainFolderPath)
    posTrainImagesPath = [posTrainFolderPath + "\\" + item for item in posImgs]

    negTrainFolderPath = ".\\INRIAPerson\\train_64x128_H96\\negNew"
    negImgs = os.listdir(negTrainFolderPath)
    negTrainImagesPath = [negTrainFolderPath + "\\" + item for item in negImgs]

    winSize = (64,128)
    blockSize = (16,16)    
    blockStride = (8,8)
    cellSize = (8,8)
    nbins = 9    
    hog = cv2.HOGDescriptor(winSize,blockSize,blockStride,cellSize,nbins) 

    trainData, trainLabels = getTrainHoGsAndLabels(hog, posTrainImagesPath, negTrainImagesPath)
    svm = LinearSVC(C = 1, loss="hinge")
    svm.fit(trainData, trainLabels)
    s = pickle.dumps(svm)
    f=open('svm.model', "wb+")
    f.write(s)
    f.close()

    posTestFolderPath = ".\\INRIAPerson\\70X134H96\\Test\\posNew"
    negTestFolderPath = ".\\INRIAPerson\\Test\\negNew"

    posTestImages = os.listdir(posTestFolderPath)
    negTestImages = os.listdir(negTestFolderPath)
    posTestImagesPath = []
    for item in posTestImages:
        posTestImagesPath.append(posTestFolderPath + "\\" + item)
    negTestImagesPath = [negTestFolderPath + "\\" + item for item in negTestImages]

    Acc = getAccuracy(svm, hog, posTestImagesPath, negTestImagesPath)
    print("Accuracy:", Acc)
    testData, testLabels = getTestHoGsAndLabels(hog, posTestImagesPath, negTestImagesPath)
    AUC = drawROCAndComputeAUC(svm, testData, testLabels)
    print("AUC:", AUC)